package com.bw.test1;

import org.apache.flink.types.Row;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.pipeline.regression.LinearRegression;
import com.alibaba.alink.pipeline.regression.LinearRegressionModel;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class LinearRegressionTest {
	@Test
	public void testLinearRegression() throws Exception {

		BatchOperator.setParallelism(1);
		// 1. 数据源
		List <Row> df = Arrays.asList(
			Row.of(2, 1, 1),
			Row.of(3, 2, 1),
			Row.of(4, 3, 2),
			Row.of(2, 4, 1),
			Row.of(2, 2, 1),
			Row.of(4, 3, 2),
			Row.of(1, 2, 1)
		);
		// 选择特征和label
		BatchOperator <?> batchData = new MemSourceBatchOp(df, "f0 int, f1 int, label int");
		String[] colnames = new String[] {"f0", "f1"};


		//  训练模型
		LinearRegression lr = new LinearRegression()
			.setFeatureCols(colnames)
			.setLabelCol("label")
			.setPredictionCol("pred")
			.enableLazyPrintModelInfo();
		LinearRegressionModel model = lr.fit(batchData);


		// 输出结果
		model.transform(batchData).print();
	}
}